{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from caffe2.python import core, workspace, test_util, dyndep, nomnigraph as ng\n",
    "from caffe2.proto import caffe2_pb2\n",
    "import pprint as pp\n",
    "import graphviz as gv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load model from a protobuf file\n",
    "Load a caffe2 model from protobuf and convert it to nomnigraph representation (https://github.com/pytorch/pytorch/tree/master/caffe2/core/nomnigraph)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Protobuf file of the model.\n",
    "MODEL_FILE = \"model.pb\"\n",
    "with open(MODEL_FILE, 'rb') as f:\n",
    "    netdef_proto = caffe2_pb2.NetDef()\n",
    "    netdef_proto.ParseFromString(f.read())\n",
    "    nnmodule = ng.NNModule(netdef_proto)\n",
    "    dfGraph = nnmodule.dataFlow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simple graph exploration\n",
    "Basic graph traversal is supported by nomnigraph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Number of operators\n",
    "print(len(dfGraph.operators))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print operator names\n",
    "for op in dfGraph.operators:\n",
    "    print(op.name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get inputs, outputs for a particular operator\n",
    "OP_NAME = \"Mul\"\n",
    "for op in dfGraph.operators:\n",
    "    if op.name == OP_NAME:\n",
    "        print(op.name)\n",
    "        print(\"Inputs\")\n",
    "        pp.pprint([tensor.name for tensor in op.inputs])\n",
    "        print(\"Outputs\")\n",
    "        pp.pprint([tensor.name for tensor in op.outputs])\n",
    "        break\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dotString = str(dfGraph)\n",
    "gvSource = gv.Source(dotString)\n",
    "gvSource"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Subgraph matching\n",
    "Nomnigraph can be used to perform subgraph pattern matching."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mg = ng.NNMatchGraph()\n",
    "matchMul = mg.createNode(ng.NeuralNetOperator(\"Mul\"), strict=True)\n",
    "matchT = mg.createNode(ng.NeuralNetData(\"*\"), strict=True)\n",
    "matchReplaceNan = mg.createNode(ng.NeuralNetOperator(\"ReplaceNaN\"))\n",
    "mg.createEdge(matchMul, matchT)\n",
    "mg.createEdge(matchT, matchReplaceNan)\n",
    "\n",
    "matches = nnmodule.match(mg)\n",
    "for match in matches:\n",
    "    # TODO: visualize subgraph\n",
    "    for node in match.nodes:\n",
    "        print(node.name)\n",
    "    break"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
